/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com>
 *
 */
#ifndef SRC_PACKETDISPATCHER_H_
#define SRC_PACKETDISPATCHER_H_

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <chrono>
#include <iomanip>
#include <pcap.h>
#include <boost/asio.hpp>
#include <boost/array.hpp>
#include <boost/bind/bind.hpp>
#include <boost/asio/signal_set.hpp>
#include <boost/version.hpp>
#include <exception>
#include <sys/resource.h>
#if defined(IS_LINUX) && defined(HAVE_NETFILTER_QUEUE)
#include <linux/netfilter.h>
#include <libnetfilter_queue/libnetfilter_queue.h>
#endif
#include "NetworkStack.h"
#include "Multiplexer.h"
#include "protocols/ethernet/EthernetProtocol.h"
#include "Protocol.h"
#include "StackLan.h"
#include "StackMobile.h"
#include "StackLanIPv6.h"
#include "StackVirtual.h"
#include "StackOpenFlow.h"
#include "StackMobileIPv6.h"
#if defined(PYTHON_BINDING) || defined(RUBY_BINDING) || defined(LUA_BINDING)
#include "Interpreter.h"
#include "TimerManager.h"
#endif
#include "EvidenceManager.h"
#include "OutputManager.h"
#include "Message.h"
#include "HTTPSession.h"
#include "Logger.h"
#include "ipset/IPRadixTree.h"
#include "Netlink.h"

#if !defined(PCAP_NETMASK_UNKNOWN)
/*
 *  This value depending on the pcap library is defined or not
 *
 */
#define PCAP_NETMASK_UNKNOWN    0xFFFFFFFF
#endif

namespace aiengine {

#define NETFILTER_QUEUE_BUFFER	4096
#define PACKET_RECVBUFSIZE	1024*256    // receive_from buffer size for a single datagram
#define PCAP_BUFFER_SIZE	4194304 // Default 4Mb for pcap buffer size
#define NETLINK_BUFFER		2048
#define BOOST_ASIO_DISABLE_EPOLL

typedef boost::asio::posix::stream_descriptor PcapStream;
typedef std::shared_ptr<PcapStream> PcapStreamPtr;

class PacketDispatcher {
public:

	enum class PacketDispatcherStatus : uint8_t {
        	RUNNING = 0,
        	STOPPED
	};

	enum class DeviceStatus : uint8_t {
		RUNNING = 0,
		UNKNOWN,
		DOWN
	};

	class Statistics {
	public:
		explicit Statistics() {
			last_total_packets_sample = 0;
			last_total_bytes_sample = 0;
			std::time(&packet_time);
			std::time(&last_packet_time);
			std::time(&first_packet_time);
		}
		virtual ~Statistics() {}

		// The variables are mutable because we change when the user prints the packetdispatcher,
		// just to avoid compute the on the packet processing and better when the user wants the info.
		mutable std::time_t packet_time;
		mutable std::time_t last_packet_time;
		mutable std::time_t first_packet_time; // For measure time of pcap files;
		mutable uint64_t last_total_packets_sample;
		mutable uint64_t last_total_bytes_sample;
	};

	explicit PacketDispatcher(const std::string &source, int buffer_size);
	explicit PacketDispatcher(const std::string &source):
		PacketDispatcher(source, PCAP_BUFFER_SIZE) {}
	explicit PacketDispatcher():PacketDispatcher("") {}

    	virtual ~PacketDispatcher();

	void open(const std::string &source);
	void open(const std::string &source, int buffer_size);
	void run(void);
	void close(void);
    	void stop(void) { io_service_.stop(); }
	void setPcapFilter(const char *filter);
	const char *getPcapFilter() const { return pcap_filter_.c_str(); }
	void status(void);
	const char *getStackName() const { return stack_name_.c_str(); }
	uint32_t getPcapFilesDuration() const { return pcap_files_duration_; }

	void setEvidences(bool value);
	bool getEvidences() const { return have_evidences_; }
	const char *getEvidencesFilename() const { return em_->filename(); }
	uint32_t getEvidencesFileSize() const { return em_->total_bytes_on_file(); }

	void statistics();
	void statistics(std::basic_ostream<char> &out) const;
	void statistics(Json &out) const;

#if defined(STAND_ALONE_TEST) || defined(TESTING)
	// Use for the tests, limits the number of packets injected
	void setMaxPackets(int packets);
	void setFromPackets(int packets);
#endif

	// Enables/Disables a HTTP server for retrieve JSON information
	void setHTTPPort(int port);
	int getHTTPPort() const;

#if defined(PYTHON_BINDING)

	// For implement the 'with' statement in python needs the methods __enter__ and __exit__
	PacketDispatcher& __enter__();
	bool __exit__(boost::python::object type, boost::python::object val, boost::python::object traceback);

	void forwardPacket(const std::string &packet, int length);
	void addTimer(PyObject *callback, int seconds);
	void removeTimer(int seconds);

	void setStack(const boost::python::object &stack);
	boost::python::object getStack() const { return pystack_; }

	const char *getStatus() const;

        // The flow have been marked as accept or drop (for external Firewall integration (Netfilter))
        bool isPacketAccepted() const { return current_packet_.isAccept(); }

	void addAuthorizedIPAddresses(const boost::python::list &items);
	boost::python::list getAuthorizedIPAddresses() const;

	void showSystemProcessInformation() const;

	// sets the listen IP for the HTTP server
	void setHTTPIPAddress(const std::string &ip);
	const char *getHTTPIPAddress() const;

	void setAuthorizationCodeToken(const std::string &value);
	const char *getAuthorizationCodeToken() const;
	const char *getRawAuthorizationCodeToken() const;

	const char *getNetworkDeviceStatus() const { return device_is_ready_ ? "up" : "down"; }
	void setNetworkDeviceStatusCallback(PyObject *callback);
	PyObject *getNetworkDeviceStatusCallback() const { return network_callback_.getCallback(); }

	void setPcapBufferSize(int32_t value);
	uint32_t getPcapBufferSize() const { return pcap_buffer_size_; }

	void setAlarmCallback(PyObject *callback);
	Callback *getAlarmCallback() { return &alarm_callback_; }
	PyObject *getPythonAlarmCallback() const { return alarm_callback_.getCallback(); }

	void setLogRestApiPythonCode(bool value) { log_rest_api_python_code_ = value; }
	bool getLogRestApiPythonCode() const { return log_rest_api_python_code_; }

#else
        void setStack(StackLan &stack);
        void setStack(StackMobile &stack);
        void setStack(StackLanIPv6 &stack);
        void setStack(StackVirtual &stack);
        void setStack(StackOpenFlow &stack);
        void setStack(StackMobileIPv6 &stack);
#endif

	uint64_t getTotalBytes(void) const { return total_bytes_; }
	uint64_t getTotalPackets(void) const { return total_packets_; }

	// This only works with real network devices
	uint64_t getTotalReceivedPackets(void) const;
	uint64_t getTotalDroppedPackets(void) const;
	uint64_t getTotalIfDroppedPackets(void) const;

	void setStack(const SharedPointer<NetworkStack> &stack);

	void setDefaultMultiplexer(MultiplexerPtr mux); // just use for the unit tests
	void setIdleFunction(std::function <void ()> idle_function) { idle_function_ = idle_function; }

	friend std::ostream& operator<< (std::ostream &out, const PacketDispatcher &pd);

#if defined(PYTHON_BINDING) || defined(RUBY_BINDING)
	void setShell(bool enable);
	bool getShell() const;

	void setLogUserCommands(bool enable);
	bool getLogUserCommands() const;
#endif

#if defined(LUA_BINDING)
	void setShell(lua_State *L, bool enable);
	bool getShell() const;
	void addTimer(lua_State* L, const char *callback, int seconds);
#endif

#if defined(RUBY_BINDING)
	void addTimer(VALUE callback, int seconds);
#endif
	void showCurrentPayloadPacket(std::basic_ostream<char>& out) const;
	void showCurrentPayloadPacket(Json &out) const;
	void showCurrentPayloadPacket() const;

	void addAuthorizedIPAddress(const std::string &addr);

private:
	void set_stack(const SharedPointer<NetworkStack> &stack);
	void start_read_ethernet_network(void);
	void start_read_raw_network(void);
#if defined(PYTHON_BINDING) || defined(RUBY_BINDING) || defined(LUA_BINDING)
	void start_read_user_input(void);
#endif
#if defined(IS_LINUX)
	void start_read_netlink(void);
	void read_netlink(boost::system::error_code error);
#if defined(HAVE_NETFILTER_QUEUE)
	void open_netfilter_queue(const int nqueue);
	static int netfilter_callback(struct nfq_q_handle *qh, struct nfgenmsg *nfmsg, struct nfq_data *nfa, void *data);
	void run_netfilter(void);
	void close_netfilter(void);
	void start_read_netfilter(void);
	void read_netfilter(boost::system::error_code error);
#endif
#endif
	void accept_http_connections();
	void read_ethernet_network(boost::system::error_code error);
	void read_raw_network(boost::system::error_code error);
	void forward_current_ethernet_packet();
	void forward_current_raw_packet();
	void restart_timer(int seconds);
        void open_device(const std::string &device, uint32_t buffer_size);
        void open_fifo_file(const std::string &fifoname);
        void close_device(void);
        void open_pcap_file(const std::string &filename);
        void close_pcap_file(void);
        void run_device(void);
        void run_pcap(void);
	void info_message(const std::string &msg);
	void error_message(const std::string &msg);
	int get_mtu_of_network_device(const std::string &name);
	void compute_packets_bytes_rate(uint64_t &packets_per_second, uint64_t &bytes_per_second) const;
	bool is_ethernet_present() const;
	int datalink_size() const;
	void bind_to_http(int port);
	bool bind_to(const std::string &ip, const std::string &device) const;
#if defined(IS_DARWIN)
	DeviceStatus network_device_status() const;
#endif
	PacketDispatcherStatus status_ = PacketDispatcherStatus::STOPPED;
	PcapStreamPtr stream_ {};
	bool pcap_file_ready_ = false; // For read packets from pcaps
	bool read_in_progress_ = false;
	bool device_is_ready_ = false; // For read packets from network devices
	bool have_evidences_ = false;
	bool pcap_is_fifo_ = false; // For read packets from fifo devices
#if defined(IS_LINUX)
	bool netfilter_is_ready_ = false; // For read packets from netfilter
#endif
#if defined(STAND_ALONE_TEST) || defined(TESTING)
	uint32_t max_packets_ = std::numeric_limits<uint32_t>::max();
	uint32_t from_packets_ = 0;
#endif
	uint64_t total_packets_ = 0;
	uint64_t total_bytes_ = 0;
	uint64_t total_tcp_server_connections_ = 0;
    	pcap_t* pcap_ = nullptr;
	boost::asio::io_service io_service_;
	boost::asio::signal_set signals_;
	Statistics stats_ {};
	struct pcap_pkthdr *packet_header_ = nullptr;
	const uint8_t *packet_ = nullptr;
	uint16_t packet_length_ = 0;
	bool io_running_ = false;
#if !defined(LUA_BINDING) && !defined(RUBY_BINDING) && !defined(JAVA_BINDING) && !defined(GO_BINDING)
	std::function <void ()> idle_function_ = [&] (void) {};
#else
	std::function <void ()> idle_function_;
#endif
	EthernetProtocolPtr eth_ = nullptr;
	Packet current_packet_ {};
	MultiplexerPtr defMux_= nullptr;
	std::string stack_name_ = "";
	std::string input_name_ = "";
	std::string pcap_filter_ = "";
	uint32_t pcap_buffer_size_ = PCAP_BUFFER_SIZE;
	uint32_t pcap_files_duration_ = 0; // Number of seconds of duration of the pcap files
	long pcap_file_header_offset_ = 0;
#if defined(IS_LINUX)
	boost::array<char, NETLINK_BUFFER> netlink_buffer_;
	SharedPointer<NetlinkSocket> nl_sock_ = nullptr;
#if defined(HAVE_NETFILTER_QUEUE)
	uint64_t total_accept_packets_ = 0;
	uint64_t total_drop_packets_ = 0;
	uint64_t total_accept_bytes_ = 0;
	uint64_t total_drop_bytes_ = 0;
	struct nfq_handle *nfq_h_ = nullptr;
	struct nfq_q_handle *nfq_qh_ = nullptr;
	int netfilter_fd_;
	uint8_t netfilter_buffer_[NETFILTER_QUEUE_BUFFER] __attribute__ ((aligned));
#endif
#endif
	SharedPointer<EvidenceManager> em_ = SharedPointer<EvidenceManager>(new EvidenceManager());
	SharedPointer<NetworkStack> current_network_stack_ = nullptr;
#if defined(PYTHON_BINDING) || defined(RUBY_BINDING) || defined(LUA_BINDING)
	SharedPointer<TimerManager> tm_ = SharedPointer<TimerManager>(new TimerManager(io_service_));
	SharedPointer<Interpreter> user_shell_ = SharedPointer<Interpreter>(new Interpreter(io_service_));
	Callback network_callback_;
	Callback alarm_callback_;
#if defined(PYTHON_BINDING)
	boost::python::object pystack_ {};
	std::string authorization_code_token_;
	bool log_rest_api_python_code_ = false;
#endif
#endif
	boost::asio::ip::tcp::socket http_socket_ { io_service_ };
	boost::asio::ip::tcp::acceptor http_acceptor_ { io_service_ };
	boost::asio::ip::address http_ip_address_ = boost::asio::ip::make_address("127.0.0.1");
	// List of IP address that can check the HTTP interface
	SharedPointer<IPRadixTree> allow_address_ = SharedPointer<IPRadixTree>(new IPRadixTree());
};

typedef std::shared_ptr<PacketDispatcher> PacketDispatcherPtr;

} // namespace aiengine

#endif  // SRC_PACKETDISPATCHER_H_
